package org.robotfish.learn.tests;

import static org.junit.Assert.*;

import org.ejml.simple.SimpleMatrix;
import org.junit.Test;
import org.robotfish.learn.EmpiricalHydroModel;

public class EmpiricalHydroModelTest {

	@Test
	public void testEHMBasic() {
		EmpiricalHydroModel ehm = new EmpiricalHydroModel(
				0.025, 1.5, 0.1, 11, 10, 0.025, 0.5, -5, 5);
		ehm.learnHydroModel(
			new SimpleMatrix( new double[][]{{0.1}, {0.2}, {-0.05}} ), // force
			new SimpleMatrix( new double[][]{{1.3}, {-3.2}, {0.1}} ), // velocity
			new SimpleMatrix( new double[][]{{105.7}, {-40.2}, {0.25}} ), // COG
			3.5, // curvature,
			1    // time step
		);
		
		assertEquals(1, ehm.getnData());
		
		ehm.learnHydroModel(
			new SimpleMatrix( new double[][]{{0.2}, {0.3}, {-0.07}} ), // force
			new SimpleMatrix( new double[][]{{2.3}, {-2.2}, {0.05}} ), // velocity
			new SimpleMatrix( new double[][]{{102.7}, {-41.2}, {0.45}} ), // COG
			4,   // curvature,
			2    // time step
		);
		assertEquals(2, ehm.getnData());

	}

	@Test
	public void testSequence() {
		assertArrayEquals(new double[]{0, 0.5, 1, 1.5}, EmpiricalHydroModel.sequence(0, 2, 4, false), 1e-5);
		assertArrayEquals(new double[]{0, 0.5, 1, 1.5, 2.0}, EmpiricalHydroModel.sequence(0, 2, 5, true), 1e-5);
		assertArrayEquals(new double[]{2.0, 1.5, 1.0, 0.5, 0}, EmpiricalHydroModel.sequence(2, 0, 5, true), 1e-5);
	}
	
	
	SimpleMatrix forceModel(SimpleMatrix velocity, SimpleMatrix cog, double k, double t) {
		SimpleMatrix out = new SimpleMatrix(3,1);
		// FT = Gauss(t) + sign(v_T) v_T^2 + Gauss(k) w^2
		out.set( 0, 2*Math.signum(velocity.get(0))*velocity.get(0)*velocity.get(0) );
		// FN = Gauss(t) + sign(v_N) v_N^2:
		out.set( 1, 2*Math.signum(velocity.get(1))*velocity.get(1)*velocity.get(1) );
		// T = Gauss(t) + sign(w) w^2 + Gauss(k) v_T^2
		out.set( 2, 2*Math.signum(velocity.get(2))*velocity.get(2)*velocity.get(2) );
		
		return out;
	}
	
	double cycle(double x, double cycleLength) {
		double c = Math.abs(x) % cycleLength;
		c = Math.min(c, cycleLength-c);
		return c;
	}
	
	@Test
	public void testFind() {
		EmpiricalHydroModel ehm = new EmpiricalHydroModel(
				0.025, 1.5, 0.1, 11, 10, 0.025, 0.5, -5, 5);
		ehm.setDt(0.025);
		ehm.setCycleLength(1.0);
		ehm.setRegularization(0.1);
		ehm.setupModels();
		int n = 100;
		for (int i=0; i<n; i++) {
			double i1 = (i+10) % 50;
			double i2 = (i+30) % 50;
			double i3 = i % 50;
			double k = ((i-50)/10.0) % 8.3;
			SimpleMatrix velocity = new SimpleMatrix( new double[][]{ {i1}, {i2}, {i3} } );
			SimpleMatrix cog = new SimpleMatrix( new double[][]{ {i}, {1.0}, {0.0} } );
			SimpleMatrix force = forceModel(velocity, cog, k, i*0.25);
			// adding points:
			ehm.learnHydroModel(force, velocity, cog, k, i);
		}
		assertEquals( n, ehm.getnData() );
		// TODO check models:
	}

}
